{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "EgiF12Hf1Dhs"
   },
   "source": [
    "**I recommend you run the first code cell of this notebook immediately, to start provisioning drake on the cloud machine, then you can leave this window open as you [read the textbook](manipulation.csail.mit.edu/pose.html).**\n",
    "\n",
    "# Notebook setup\n",
    "\n",
    "The following cell will:\n",
    "- on Colab (only), install Drake to `/opt/drake`, install Drake's prerequisites via `apt`, and add pydrake to `sys.path`.  This will take approximately two minutes on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  If you navigate between notebooks using Colab's \"File->Open\" menu, then you can avoid provisioning a separate machine for each notebook.\n",
    "- launch a server for our 3D visualizer (MeshCat) that will be used for the remainder of this notebook.\n",
    "\n",
    "You will need to rerun this cell if you restart the kernel, but it should be fast because the machine will already have drake installed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "eeMrMI0-1Dhu"
   },
   "outputs": [],
   "source": [
    "import importlib\n",
    "import sys\n",
    "from urllib.request import urlretrieve\n",
    "\n",
    "if 'google.colab' in sys.modules and importlib.util.find_spec('manipulation') is None:\n",
    "    urlretrieve(f\"http://manipulation.csail.mit.edu/scripts/setup/setup_manipulation_colab.py\",\n",
    "                \"setup_manipulation_colab.py\")\n",
    "    from setup_manipulation_colab import setup_manipulation\n",
    "    setup_manipulation(manipulation_sha='master', drake_version='20200915', drake_build='nightly')\n",
    "\n",
    "# Determine if this notebook is currently running as a notebook or a unit test.\n",
    "from IPython import get_ipython\n",
    "running_as_notebook = get_ipython() and hasattr(get_ipython(), 'kernel')\n",
    "\n",
    "# Setup rendering (with xvfb), if necessary:\n",
    "import os\n",
    "if 'google.colab' in sys.modules and os.getenv(\"DISPLAY\") is None:\n",
    "    from pyvirtualdisplay import Display\n",
    "    display = Display(visible=0, size=(1400, 900))\n",
    "    display.start()\n",
    "\n",
    "# Start a single meshcat server instance to use for the remainder of this notebook.\n",
    "from meshcat.servers.zmqserver import start_zmq_server_as_subprocess\n",
    "server_args = []\n",
    "if 'google.colab' in sys.modules:\n",
    "    server_args = ['--ngrok_http_tunnel']\n",
    "proc, zmq_url, web_url = start_zmq_server_as_subprocess(server_args=server_args)\n",
    "\n",
    "# Let's do all of our imports here, too.\n",
    "import numpy as np\n",
    "import matplotlib.animation as animation\n",
    "import matplotlib.pyplot as plt, mpld3\n",
    "import open3d\n",
    "import pydot\n",
    "from IPython.display import display, SVG, HTML\n",
    "\n",
    "from pydrake.all import (\n",
    "    AddMultibodyPlantSceneGraph, BaseField, Box, ConnectMeshcatVisualizer, CsdpSolver, DepthCameraProperties, \n",
    "    DepthImageToPointCloud, DiagramBuilder, FindResourceOrThrow, FixedOffsetFrame, \n",
    "    ge, GeometryInstance, MakeRenderEngineVtk, MakePhongIllustrationProperties, \n",
    "    MathematicalProgram, MeshcatPointCloudVisualizer, Parser, RandomGenerator, RenderEngineVtkParams, RgbdSensor, RigidTransform, \n",
    "    RollPitchYaw, RotationMatrix, SpatialInertia, Solve, set_log_level, UniformlyRandomRotationMatrix, world_model_instance)\n",
    "from pydrake.examples.manipulation_station import ManipulationStation\n",
    "\n",
    "from manipulation.utils import FindResource\n",
    "\n",
    "if running_as_notebook:\n",
    "    mpld3.enable_notebook()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "7q0A14bAilIX"
   },
   "source": [
    "# Simulating an RGB-D camera\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ILYLouFTjv6e"
   },
   "outputs": [],
   "source": [
    "def DepthCameraDemoSystem():\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    # Create the physics engine + scene graph.\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)\n",
    "    # Add a single object into it.\n",
    "    X_Mustard = RigidTransform(RollPitchYaw(-np.pi/2., 0, -np.pi/2.), [0, 0, 0.09515])\n",
    "    parser = Parser(plant)\n",
    "    mustard = parser.AddModelFromFile(FindResourceOrThrow(\n",
    "        \"drake/manipulation/models/ycb/sdf/006_mustard_bottle.sdf\"))\n",
    "    plant.WeldFrames(plant.world_frame(), \n",
    "                     plant.GetFrameByName(\"base_link_mustard\", mustard), \n",
    "                     X_Mustard)\n",
    "    # Add a rendering engine\n",
    "    renderer = \"my_renderer\"\n",
    "    scene_graph.AddRenderer(renderer,\n",
    "                            MakeRenderEngineVtk(RenderEngineVtkParams()))\n",
    "\n",
    "    # Add a box for the camera in the environment.\n",
    "    X_Camera = RigidTransform(\n",
    "        RollPitchYaw(0, -0.2, 0.2).ToRotationMatrix().multiply(\n",
    "            RollPitchYaw(-np.pi/2.0, 0, np.pi/2.0).ToRotationMatrix()),\n",
    "        [.5, .1, .2])\n",
    "    camera_instance = parser.AddModelFromFile(FindResource(\"models/camera_box.sdf\"))\n",
    "    camera = plant.GetBodyByName(\"base\", camera_instance)    \n",
    "    plant.WeldFrames(plant.world_frame(), camera.body_frame(), X_Camera)\n",
    "\n",
    "    plant.Finalize()\n",
    "\n",
    "    frames_to_draw = {\"cameras\": \"camera\"}\n",
    "    meshcat = ConnectMeshcatVisualizer(builder, scene_graph, zmq_url=zmq_url, \n",
    "                                       delete_prefix_on_load=False,\n",
    "                                       frames_to_draw=frames_to_draw)\n",
    "    meshcat.load()\n",
    "\n",
    "    properties = DepthCameraProperties(width=640,\n",
    "                                       height=480,\n",
    "                                       fov_y=np.pi / 4.0,\n",
    "                                       renderer_name=renderer,\n",
    "                                       z_near=0.1,\n",
    "                                       z_far=10.0)\n",
    "    camera = builder.AddSystem(\n",
    "        RgbdSensor(parent_id=plant.GetBodyFrameIdOrThrow(camera.index()),\n",
    "                   X_PB=RigidTransform(),\n",
    "                   properties=properties,\n",
    "                   show_window=False))\n",
    "    camera.set_name(\"rgbd_sensor\")\n",
    "    builder.Connect(scene_graph.get_query_output_port(),\n",
    "                    camera.query_object_input_port())\n",
    "\n",
    "    # Export the camera outputs\n",
    "    builder.ExportOutput(camera.color_image_output_port(), \"color_image\")\n",
    "    builder.ExportOutput(camera.depth_image_32F_output_port(), \"depth_image\")\n",
    "\n",
    "    # Add a system to convert the camera output into a point cloud\n",
    "    to_point_cloud = builder.AddSystem(\n",
    "        DepthImageToPointCloud(camera_info=camera.depth_camera_info(),\n",
    "                               fields=BaseField.kXYZs | BaseField.kRGBs))\n",
    "    builder.Connect(camera.depth_image_32F_output_port(),\n",
    "                    to_point_cloud.depth_image_input_port())\n",
    "    builder.Connect(camera.color_image_output_port(),\n",
    "                    to_point_cloud.color_image_input_port())\n",
    "\n",
    "    # Send the point cloud to meshcat for visualization, too.\n",
    "    meshcat_pointcloud = builder.AddSystem(MeshcatPointCloudVisualizer(meshcat, X_WP=X_Camera))\n",
    "    builder.Connect(to_point_cloud.point_cloud_output_port(), meshcat_pointcloud.get_input_port())\n",
    "\n",
    "    # Export the point cloud output.\n",
    "    builder.ExportOutput(to_point_cloud.point_cloud_output_port(),\n",
    "                         \"point_cloud\")\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    diagram.set_name(\"depth_camera_demo_system\")\n",
    "    return diagram\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "WCb1f0DmMUay",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def plot_camera_images():\n",
    "    system = DepthCameraDemoSystem()\n",
    "\n",
    "    # Evaluate the camera output ports to get the images.\n",
    "    context = system.CreateDefaultContext()\n",
    "    system.Publish(context)\n",
    "    color_image = system.GetOutputPort(\"color_image\").Eval(context)\n",
    "    depth_image = system.GetOutputPort(\"depth_image\").Eval(context)\n",
    "\n",
    "    # Plot the two images.\n",
    "    plt.subplot(121)\n",
    "    plt.imshow(color_image.data)\n",
    "    plt.title('Color image')\n",
    "    plt.subplot(122)\n",
    "    plt.imshow(np.squeeze(depth_image.data))\n",
    "    plt.title('Depth image')\n",
    "    #mpld3.display()\n",
    "    plt.show()\n",
    "\n",
    "plot_camera_images()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Wya-_6_3MUa1",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def draw_diagram():\n",
    "    system = DepthCameraDemoSystem()\n",
    "    display(SVG(pydot.graph_from_dot_data(system.GetGraphvizString(max_depth=1))[0].create_svg()))\n",
    "\n",
    "draw_diagram()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "mFNDRsZ1MUa4",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def plot_manipulation_station_camera_images():\n",
    "    station = ManipulationStation()\n",
    "    station.SetupManipulationClassStation()\n",
    "    # station.SetupClutterClearingStation()\n",
    "    station.Finalize()\n",
    "    context = station.CreateDefaultContext()\n",
    "\n",
    "    camera_names = station.get_camera_names()\n",
    "    index = 1\n",
    "    for name in camera_names:\n",
    "        color_image = station.GetOutputPort(\"camera_\" + name +\n",
    "                                            \"_rgb_image\").Eval(context)\n",
    "        depth_image = station.GetOutputPort(\"camera_\" + name +\n",
    "                                            \"_depth_image\").Eval(context)\n",
    "\n",
    "        plt.subplot(len(camera_names), 2, index)\n",
    "        plt.imshow(color_image.data)\n",
    "        index += 1\n",
    "        plt.title('Color image')\n",
    "        plt.subplot(len(camera_names), 2, index)\n",
    "        plt.imshow(np.squeeze(depth_image.data))\n",
    "        index += 1\n",
    "        plt.title('Depth image')\n",
    "\n",
    "    plt.show()\n",
    "\n",
    "plot_manipulation_station_camera_images()    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "9QDyRMCyb_gL"
   },
   "source": [
    "# Point cloud registration with known correspondences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "vU_zEQJ_b1mv",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def MakeRandomObjectModelAndScenePoints(\n",
    "    num_model_points=20, \n",
    "    noise_std=0, \n",
    "    num_outliers=0, \n",
    "    yaw_O=None,\n",
    "    p_O=None, \n",
    "    num_viewable_points=None, \n",
    "    seed=None):\n",
    "    \"\"\" Returns p_Om, p_s \"\"\"\n",
    "    random_state = np.random.RandomState(seed)\n",
    "\n",
    "    # Make a random set of points to define our object in the x,y plane\n",
    "    theta = np.arange(0, 2.0*np.pi, 2.0*np.pi/num_model_points)\n",
    "    l = 1.0 + 0.5*np.sin(2.0*theta) + 0.4*random_state.rand(1, num_model_points)\n",
    "    p_Om = np.vstack((l * np.sin(theta), l * np.cos(theta), 0 * l))\n",
    "\n",
    "    # Make a random object pose if one is not specified, and apply it to get the scene points.\n",
    "    if p_O is None:\n",
    "        p_O = [2.0*random_state.rand(), 2.0*random_state.rand(), 0.0]\n",
    "    if len(p_O) == 2:\n",
    "        p_O.append(0.0)\n",
    "    if yaw_O is None:\n",
    "        yaw_O = 0.5*random_state.random()\n",
    "    X_O = RigidTransform(RotationMatrix.MakeZRotation(yaw_O), p_O)\n",
    "    if num_viewable_points is None:\n",
    "        num_viewable_points = num_model_points\n",
    "    assert num_viewable_points <= num_model_points\n",
    "    p_s = X_O.multiply(p_Om[:,:num_viewable_points])\n",
    "    p_s[:2, :]  += random_state.normal(scale=noise_std, size=(2, num_viewable_points))\n",
    "    if num_outliers:\n",
    "        outliers = random_state.uniform(low=-1.5, high=3.5, size=(3, num_outliers))\n",
    "        outliers[2,:] = 0\n",
    "        p_s = np.hstack((p_s, outliers))\n",
    "\n",
    "    return p_Om, p_s, X_O\n",
    "\n",
    "def MakeRectangleModelAndScenePoints(\n",
    "    num_points_per_side=7,\n",
    "    noise_std=0, \n",
    "    num_outliers=0, \n",
    "    yaw_O=None,\n",
    "    p_O=None, \n",
    "    num_viewable_points=None, \n",
    "    seed=None):\n",
    "    random_state = np.random.RandomState(seed)\n",
    "    if p_O is None:\n",
    "        p_O = [2.0*random_state.rand(), 2.0*random_state.rand(), 0.0]\n",
    "    if len(p_O) == 2:\n",
    "        p_O.append(0.0)\n",
    "    if yaw_O is None:\n",
    "        yaw_O = 0.5*random_state.random()\n",
    "    X_O = RigidTransform(RotationMatrix.MakeZRotation(yaw_O), p_O)\n",
    "    if num_viewable_points is None:\n",
    "        num_viewable_points = 4*num_points_per_side\n",
    "    \n",
    "    x = np.arange(-1, 1, 2/num_points_per_side)\n",
    "    half_width = 2\n",
    "    half_height = 1\n",
    "    top = np.vstack((half_width*x, half_height + 0*x))\n",
    "    right = np.vstack((half_width + 0*x, -half_height*x))\n",
    "    bottom = np.vstack((-half_width*x, -half_height + 0*x))\n",
    "    left = np.vstack((-half_width + 0*x, half_height*x))\n",
    "    p_Om = np.vstack((np.hstack((top, right, bottom, left)), np.zeros((1, 4*num_points_per_side))))\n",
    "    p_s = X_O.multiply(p_Om[:,:num_viewable_points])\n",
    "    p_s[:2, :]  += random_state.normal(scale=noise_std, size=(2, num_viewable_points))\n",
    "    if num_outliers:\n",
    "        outliers = random_state.uniform(low=-1.5, high=3.5, size=(3, num_outliers))\n",
    "        outliers[2,:] = 0\n",
    "        p_s = np.hstack((p_s, outliers))\n",
    "\n",
    "    return p_Om, p_s, X_O\n",
    "\n",
    "\n",
    "def PlotEstimate(p_Om, p_s, Xhat_O=RigidTransform(), chat=None, X_O=None, ax=None):\n",
    "    p_m = Xhat_O.multiply(p_Om)\n",
    "    if ax is None:\n",
    "        ax = plt.subplot()\n",
    "    Nm = p_Om.shape[1]\n",
    "    artists = ax.plot(p_m[0, :], p_m[1, :], 'bo')\n",
    "    artists += ax.fill(p_m[0, :], p_m[1, :], 'lightblue', alpha=0.5)\n",
    "    artists += ax.plot(p_s[0, :], p_s[1, :], 'ro')\n",
    "    if chat is not None:\n",
    "        artists += ax.plot(np.vstack((p_m[0, chat], p_s[0, :])), np.vstack((p_m[1, chat], p_s[1, :])), 'g--')\n",
    "    if X_O:\n",
    "        p_s = X_O.multiply(p_Om)\n",
    "    artists += ax.fill(p_s[0, :Nm], p_s[1, :Nm], 'lightsalmon')\n",
    "    ax.axis('equal')\n",
    "    return artists\n",
    "\n",
    "def PrintResults(X_O, Xhat_O):\n",
    "    p = X_O.translation()\n",
    "    aa = X_O.rotation().ToAngleAxis()\n",
    "    print(f\"True position: {p}\")\n",
    "    print(f\"True orientation: {aa}\")\n",
    "    p = Xhat_O.translation()\n",
    "    aa = Xhat_O.rotation().ToAngleAxis()\n",
    "    print(f\"Estimated position: {p}\")\n",
    "    print(f\"Estimated orientation: {aa}\")\n",
    "\n",
    "def PoseEstimationGivenCorrespondences(p_Om, p_s, chat):\n",
    "    \"\"\" Returns optimal X_O given the correspondences \"\"\"\n",
    "    # Apply correspondences, and transpose data to support numpy broadcasting\n",
    "    p_Omc = p_Om[:, chat].T\n",
    "    p_s = p_s.T\n",
    "\n",
    "    # Calculate the central points\n",
    "    p_Ombar = p_Omc.mean(axis=0)\n",
    "    p_sbar = p_s.mean(axis=0)\n",
    "\n",
    "    # Calculate the \"error\" terms, and form the data matrix\n",
    "    merr = p_Omc - p_Ombar\n",
    "    serr = p_s - p_sbar\n",
    "    W = np.matmul(serr.T, merr)\n",
    "\n",
    "    # Compute R\n",
    "    U, Sigma, Vt = np.linalg.svd(W)\n",
    "    R = np.matmul(U, Vt)\n",
    "    if np.linalg.det(R) < 0:\n",
    "       print(\"fixing improper rotation\")\n",
    "       Vt[-1, :] *= -1\n",
    "       R = np.matmul(U, Vt)\n",
    "\n",
    "    # Compute p\n",
    "    p = p_sbar - np.matmul(R, p_Ombar)\n",
    "\n",
    "    return RigidTransform(RotationMatrix(R), p)\n",
    "\n",
    "\n",
    "p_Om, p_s, X_O = MakeRandomObjectModelAndScenePoints(num_model_points=20)\n",
    "#p_Om, p_s, X_O = MakeRectangleModelAndScenePoints()\n",
    "Xhat = RigidTransform()\n",
    "c = range(p_Om.shape[1])  # perfect, known correspondences\n",
    "fig, ax = plt.subplots(1, 2)\n",
    "PlotEstimate(p_Om, p_s, Xhat, c, ax=ax[0])\n",
    "Xhat = PoseEstimationGivenCorrespondences(p_Om, p_s, c)\n",
    "ax[1].set_xlim(ax[0].get_xlim())\n",
    "ax[1].set_ylim(ax[0].get_ylim());\n",
    "PlotEstimate(p_Om, p_s, Xhat, c, ax=ax[1])\n",
    "ax[0].set_title('Original Data')\n",
    "ax[1].set_title('After Registration')\n",
    "PrintResults(X_O, Xhat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "AHfxMwrvb1mz"
   },
   "source": [
    "# Iterative Closest Point (ICP)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "N2cYjTpub1m0",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def FindClosestPoints(point_cloud_A, point_cloud_B):\n",
    "    \"\"\"\n",
    "    Finds the nearest (Euclidean) neighbor in point_cloud_B for each\n",
    "    point in point_cloud_A.\n",
    "    @param point_cloud_A A 3xN numpy array of points.\n",
    "    @param point_cloud_B A 3xN numpy array of points.\n",
    "    @return indices An (N, ) numpy array of the indices in point_cloud_B of each\n",
    "        point_cloud_A point's nearest neighbor.\n",
    "    \"\"\"\n",
    "    indices = np.empty(point_cloud_A.shape[1], dtype=int)\n",
    "\n",
    "    kdtree = open3d.geometry.KDTreeFlann(point_cloud_B)\n",
    "    for i in range(point_cloud_A.shape[1]):\n",
    "        nn = kdtree.search_knn_vector_3d(point_cloud_A[:,i], 1)\n",
    "        indices[i] = nn[1][0]\n",
    "\n",
    "    return indices\n",
    "\n",
    "def IterativeClosestPoint(p_Om, p_s, X_O=None, animate=True):\n",
    "    Xhat = RigidTransform()\n",
    "    Nm = p_s.shape[1]\n",
    "    chat_previous = np.zeros(Nm)\n",
    "\n",
    "    fig, ax = plt.subplots()\n",
    "    frames = []\n",
    "    frames.append(PlotEstimate(p_Om=p_Om, p_s=p_s, Xhat_O=Xhat, chat=None, X_O=X_O, ax=ax))\n",
    "\n",
    "    while True:\n",
    "        chat = FindClosestPoints(p_s, Xhat.multiply(p_Om))\n",
    "        if np.array_equal(chat, chat_previous):\n",
    "            # Then I've converged.\n",
    "            break\n",
    "        chat_previous = chat\n",
    "        frames.append(PlotEstimate(p_Om=p_Om, p_s=p_s, Xhat_O=Xhat, chat=chat, X_O=X_O, ax=ax))\n",
    "        Xhat = PoseEstimationGivenCorrespondences(p_Om, p_s, chat)\n",
    "        frames.append(PlotEstimate(p_Om=p_Om, p_s=p_s, Xhat_O=Xhat, chat=None, X_O=X_O, ax=ax))\n",
    "\n",
    "    ani = animation.ArtistAnimation(fig, frames, interval=400, repeat=False)\n",
    "\n",
    "    display(HTML(ani.to_jshtml()))\n",
    "    plt.close()\n",
    "\n",
    "    if X_O:\n",
    "        PrintResults(X_O, Xhat)\n",
    "\n",
    "    return Xhat, chat\n",
    "    \n",
    "\n",
    "p_Om, p_s, X_O = MakeRandomObjectModelAndScenePoints(num_model_points=20)\n",
    "IterativeClosestPoint(p_Om, p_s, X_O);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "WLMLxAJJb1m2"
   },
   "source": [
    "Try increasing the standard deviation on yaw in the example above.  At some point, the performance can get pretty poor!\n",
    "\n",
    "# ICP with messy point clouds\n",
    "\n",
    "Try changing the amount of noise, the number of outliers, and/or the partial views.  There are not particularly good theorems here, but I hope that a little bit of play will get you a lot of intuition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#p_Om, p_s, X_O = MakeRandomObjectModelAndScenePoints(\n",
    "p_Om, p_s, X_O = MakeRectangleModelAndScenePoints(\n",
    "    yaw_O=0.1,\n",
    "#    noise_std=0.2,\n",
    "#    num_outliers=3,\n",
    "    num_viewable_points=14)\n",
    "IterativeClosestPoint(p_Om, p_s, X_O);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Non-penetration of a half-plane with convex optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def ConstrainedPoseEstimationGivenCorrespondences(p_Om, p_s, chat):\n",
    "    \"\"\" This version adds a non-penetration constraint (x,y >= 0) \"\"\"\n",
    "\n",
    "    p_Omc = p_Om[:2, chat]\n",
    "    p_s = p_s[:2, :]\n",
    "    Ns = p_s.shape[1]\n",
    "\n",
    "    prog = MathematicalProgram()\n",
    "    [a,b] = prog.NewContinuousVariables(2)\n",
    "    # We use the slack variable as an upper bound on the cost of each point to make the objective linear.\n",
    "    slack = prog.NewContinuousVariables(Ns)\n",
    "    p = prog.NewContinuousVariables(2)\n",
    "    prog.AddBoundingBoxConstraint(0,1,[a,b])  # This makes Csdp happier\n",
    "    R = np.array([[a, -b],[b, a]])\n",
    "    prog.AddLorentzConeConstraint([1.0, a, b])\n",
    "\n",
    "    # Note: Could do this more efficiently, exploiting trace.  But I'm keeping it simpler here.\n",
    "    prog.AddCost(np.sum(slack))\n",
    "    for i in range(Ns):\n",
    "        c = p + np.matmul(R,p_Omc[:,i]) - p_s[:,i]\n",
    "        # forall i, slack[i]^2 >= |c|^2 \n",
    "        prog.AddLorentzConeConstraint([slack[i], c[0], c[1]])\n",
    "        # forall i, p + R*mi >= 0.  \n",
    "        prog.AddConstraint(ge(p + np.matmul(R, p_Omc[:,i]), [0, 0]))\n",
    "    \n",
    "    result = CsdpSolver().Solve(prog)\n",
    "    \n",
    "    [a,b] = result.GetSolution([a,b])\n",
    "    Rsol = np.array([[a, -b, 0],[b, a, 0], [0,0,1]])\n",
    "    psol = np.zeros(3)\n",
    "    psol[:2] = result.GetSolution(p)\n",
    "\n",
    "    return RigidTransform(RotationMatrix(Rsol), psol)\n",
    "\n",
    "p_Om, p_s, X_O = MakeRectangleModelAndScenePoints(\n",
    "    yaw_O=0.2,\n",
    "    p_O = [1.5, 1.2],\n",
    ")\n",
    "c = range(p_Om.shape[1])  # perfect, known correspondences\n",
    "Xhat_O = ConstrainedPoseEstimationGivenCorrespondences(p_Om, p_s, c)\n",
    "PlotEstimate(p_Om=p_Om, p_s=p_s, Xhat_O=Xhat_O, chat=c, X_O=X_O)\n",
    "PrintResults(X_O, Xhat_O)\n",
    "plt.gca().plot([0,0], [0, 2.5], 'g-', linewidth=3)\n",
    "plt.gca().plot([0,4], [0, 0], 'g-', linewidth=3);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Robotic Manipulation - Geometric Pose Estimation.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "name": "Python 3.6.9 64-bit",
   "display_name": "Python 3.6.9 64-bit",
   "metadata": {
    "interpreter": {
     "hash": "30295c5bec572e859485b1ffa5e89b8b3e2022ef6e3e739c1ac40f143a557caf"
    }
   }
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9-final"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}